"""
分割分析服务

这是jinxiang_ZZ模块的主要接口，封装SegmentationAnalyzer，提供简化的文件夹分析功能。
"""

import os
import sys
import time
import numpy as np
from typing import Dict, Any, List, Optional
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from devdeploy.src import SlideSegmentationAnalyzer, get_optimal_device, prepare_output_directory


class SegmentationAnalyzerService:
    """
    分割分析服务
    
    这是jinxiang_ZZ模块的主要接口，提供简化的文件夹分析功能。
    """
    
    def __init__(self, model_path: str, device: Optional[str] = None):
        """
        初始化分割分析服务
        
        Args:
            model_path (str): 模型文件路径
            device (str, optional): 推理设备，如果为None则自动选择
        """
        self.model_path = model_path
        self.device = device or get_optimal_device()
        
        # 初始化分割分析器
        try:
            self.analyzer = SlideSegmentationAnalyzer(
                model_path=self.model_path,
                device=self.device
            )
            print(f"分割分析服务初始化成功，使用设备: {self.device}")
        except Exception as e:
            raise RuntimeError(f"初始化分割分析服务失败: {str(e)}")
    
    def analyze_folder(self, input_folder: str, output_folder: str, 
                      target_classes: List[int] = [1, 2]) -> Dict[str, Any]:
        """
        分析文件夹中的所有图像
        
        Args:
            input_folder (str): 输入文件夹路径
            output_folder (str): 输出文件夹路径
            target_classes (List[int]): 要分析的类别，默认为[1, 2]
            
        Returns:
            Dict[str, Any]: 分析结果
        """
        result = {
            'success': False,
            'reuslts': [],
            'message': '',
            'total_images': 0,
            'processed_images': 0,
            'failed_images': 0
        }
        
        try:
            # 验证输入文件夹
            if not os.path.exists(input_folder):
                result['message'] = f"输入文件夹不存在: {input_folder}"
                return result
            
            if not os.path.isdir(input_folder):
                result['message'] = f"路径不是文件夹: {input_folder}"
                return result
            
            # 使用prepare_output_directory确保输出文件夹正确创建
            if not prepare_output_directory(output_folder, clear_existing=True):
                result['message'] = f"无法创建输出文件夹: {output_folder}"
                return result
            
            # 创建result子文件夹用于保存图像结果
            result_folder = output_folder
            
            # 获取图像文件列表
            image_files = self._get_image_files(input_folder)
            result['total_images'] = len(image_files)
            
            if len(image_files) == 0:
                result['message'] = f"文件夹中没有找到支持的图像文件: {input_folder}"
                return result
            
            print(f"开始分析 {len(image_files)} 个图像文件...")
            
            # 存储所有分析结果用于生成报告
            all_results = []

            # 逐个处理图像
            for i, image_path in enumerate(image_files):
                try:
                    print(f"处理图像 {i+1}/{len(image_files)}: {os.path.basename(image_path)}")
                    
                    # 分析单张图像（计时）
                    start_time = time.time()
                    analysis_result = self.analyzer.analyze_single(image_path, target_classes,900,800,0.1,0,"max")
                    elapsed_analyze = time.time() - start_time
                    print(f"  - 单图像分析耗时: {elapsed_analyze:.3f} 秒")

                    # 获取文件名（不含扩展名）
                    file_name = Path(image_path).stem
                    
                    # 创建可视化结果（计时）
                    vis_start_time = time.time()
                    visualization_path = os.path.join(result_folder, f"{file_name}_visualization.png")
                    self._create_visualization(image_path, analysis_result, visualization_path)
                    elapsed_vis = time.time() - vis_start_time
                    print(f"  - 可视化生成耗时: {elapsed_vis:.3f} 秒")
                    
                    # 记录结果用于报告，包含类别名称信息（计时）
                    record_start_time = time.time()
                    all_results.append({
                        'file_name': file_name,
                        'class_ratios': analysis_result['class_ratios'],
                        'class_names': analysis_result.get('class_names', {}),
                        'image_shape': analysis_result['image_shape']
                    })
                    elapsed_record = time.time() - record_start_time
                    print(f"  - 结果记录耗时: {elapsed_record:.3f} 秒")
                    
                    result['processed_images'] += 1
                    
                except Exception as e:
                    print(f"处理图像失败 {image_path}: {str(e)}")
                    result['failed_images'] += 1
            
            # 生成汇总报告到根目录
            self._generate_summary_report(all_results, output_folder, target_classes)
            
            result['success'] = True
            result['results'] = all_results
            result['message'] = f"分析完成！成功处理 {result['processed_images']} 个文件，失败 {result['failed_images']} 个文件"
            
        except Exception as e:
            result['message'] = f"分析文件夹时发生错误: {str(e)}"
        
        return result
    
    def _get_image_files(self, folder_path: str) -> List[str]:
        """
        获取文件夹中的所有图像文件
        
        Args:
            folder_path (str): 文件夹路径
            
        Returns:
            List[str]: 图像文件路径列表
        """
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        image_files = []
        
        for file_path in Path(folder_path).iterdir():
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                image_files.append(str(file_path))
        
        return sorted(image_files)
    
    
    def _generate_summary_report(self, all_results: List[Dict], output_folder: str, target_classes: List[int]):
        """
        生成汇总报告到根目录的result.txt文件
        
        Args:
            all_results (List[Dict]): 所有分析结果
            output_folder (str): 输出文件夹
            target_classes (List[int]): 目标类别
        """
        try:
            summary_path = os.path.join(output_folder, 'result.txt')
            
            with open(summary_path, 'w', encoding='utf-8') as f:
                f.write("金相组织分析汇总报告\n")
                f.write("=" * 50 + "\n")
                f.write(f"分析时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"模型路径: {self.model_path}\n")
                f.write(f"目标类别: {target_classes}\n")
                f.write(f"总文件数: {len(all_results)}\n\n")
                
                if all_results:
                    # 获取类别名称映射（从第一个结果中获取）
                    class_names = all_results[0].get('class_names', {})
                    
                    # 计算平均占比
                    avg_ratios = {}
                    for class_id in target_classes:
                        ratios = [r['class_ratios'].get(class_id, 0) for r in all_results]
                        avg_ratios[class_id] = sum(ratios) / len(ratios)
                    
                    f.write("整体统计:\n")
                    for class_id, avg_ratio in avg_ratios.items():
                        class_name = class_names.get(class_id, f"类别{class_id}")
                        f.write(f"{class_name} (ID:{class_id}) 平均占比: {avg_ratio:.2f}%\n")
                    f.write("\n")
                    
                    f.write("详细结果:\n")
                    for r in all_results:
                        f.write(f"{r['file_name']} - ")
                        for class_id in target_classes:
                            ratio = r['class_ratios'].get(class_id, 0)
                            class_name = class_names.get(class_id, f"类别{class_id}")
                            f.write(f"{class_name}: {ratio:.2f}% ")
                        f.write("\n")
            
            print(f"汇总报告已保存: {summary_path}")
            
        except Exception as e:
            print(f"生成汇总报告失败: {str(e)}")
    
    def _create_visualization(self, original_image_path: str, result: Dict[str, Any], output_path: str) -> bool:
        """
        创建简化的可视化结果，只保存overlap图像以提高性能
        
        Args:
            original_image_path: 原始图像路径
            result: 分析结果
            output_path: 输出路径
            
        Returns:
            bool: 是否保存成功
        """
        try:
            # 加载原始图像
            img = Image.open(original_image_path)
            
            # 获取结果数据
            class_ratios = result['class_ratios']
            class_masks = result['class_masks']
            class_names = result.get('class_names', {})
            
            # 创建overlap图像（所有类别叠加）
            overlap_img = self._create_overlap_visualization(img, class_masks, class_ratios, class_names)
            
            if overlap_img is None:
                print(f"无法创建overlap图像: {output_path}")
                return False
            
            # 创建单个子图显示overlap图像
            fig, ax = plt.subplots(1, 1, figsize=(10, 8))
            
            # 显示overlap图像
            ax.imshow(overlap_img)
            
            # 动态生成标题（使用类别名称）
            available_classes = sorted([cid for cid in class_masks.keys() if cid > 0])
            if available_classes:
                title_parts = [class_names.get(cid, f"类别{cid}") for cid in available_classes]
                ax.set_title(' + '.join(title_parts) + ' Overlap')
            else:
                ax.set_title('Class Overlap')
            ax.axis('off')
            
            # 动态生成统计信息（使用类别名称）
            stats_lines = []
            for class_id in sorted(class_ratios.keys()):
                if class_id > 0:  # 跳过背景类别0
                    class_name = class_names.get(class_id, f"类别{class_id}")
                    ratio = class_ratios.get(class_id, 0)
                    stats_lines.append(f"{class_name}: {ratio:.2f}%")
            
            stats_text = "\n".join(stats_lines) if stats_lines else "无类别数据"
            ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, 
                   fontsize=14, fontweight='bold', color='white',
                   bbox=dict(boxstyle="round,pad=0.3", facecolor="black", alpha=0.7),
                   verticalalignment='top')
            
            # 保存图像
            plt.tight_layout()
            plt.savefig(output_path, dpi=150, bbox_inches='tight')
            plt.close()
            
            return True
            
        except Exception as e:
            print(f"创建可视化失败 {output_path}: {str(e)}")
            return False
    
    def _create_overlap_visualization(self, original_img: Image.Image, class_masks: Dict, class_ratios: Dict, class_names: Dict = None) -> Optional[np.ndarray]:
        """
        创建所有类别的overlap可视化
        
        Args:
            original_img: 原始图像
            class_masks: 类别掩码字典
            class_ratios: 类别占比字典
            class_names: 类别名称字典，可选
            
        Returns:
            np.ndarray: overlap图像
        """
        try:
            if class_names is None:
                class_names = {}
            
            # 将PIL图像转换为numpy数组
            if isinstance(original_img, Image.Image):
                img_array = np.array(original_img)
            else:
                img_array = original_img
            
            # 获取所有非背景类别（类别ID > 0）
            available_classes = sorted([cid for cid in class_masks.keys() if cid > 0])
            
            if not available_classes:
                return None
            
            # 定义颜色列表（RGB格式）
            colors = [
                [0, 255, 0],      # 绿色
                [255, 0, 0],      # 红色
                [0, 0, 255],      # 蓝色
                [255, 255, 0],    # 黄色
                [255, 0, 255],    # 洋红色
                [0, 255, 255],    # 青色
                [255, 128, 0],    # 橙色
                [128, 0, 255],    # 紫色
                [255, 192, 203],  # 粉色
                [128, 128, 128],  # 灰色
            ]
            
            # 创建总的overlay
            total_overlay = np.zeros_like(img_array)
            
            # 为每个类别创建overlay
            for idx, class_id in enumerate(available_classes):
                mask = class_masks[class_id]
                color = colors[idx % len(colors)]  # 循环使用颜色
                
                # 应用颜色到该类别区域
                mask_bool = mask > 0
                total_overlay[mask_bool] = color
            
            # 处理重叠区域：如果有多个类别重叠，使用混合颜色
            if len(available_classes) > 1:
                # 计算重叠区域
                overlap_mask = np.zeros(img_array.shape[:2], dtype=bool)
                for class_id in available_classes:
                    overlap_mask |= (class_masks[class_id] > 0)
                
                # 对于重叠区域，使用白色或混合色
                overlap_count = np.zeros(img_array.shape[:2], dtype=int)
                for class_id in available_classes:
                    overlap_count += (class_masks[class_id] > 0).astype(int)
                
                # 重叠区域（2个或更多类别）使用白色
                multi_overlap = overlap_count >= 2
                total_overlay[multi_overlap] = [255, 255, 255]  # 白色表示重叠
            
            # 将overlay与原图混合
            alpha = 0.6  # 透明度
            overlap_img = (1 - alpha) * img_array + alpha * total_overlay
            overlap_img = np.clip(overlap_img, 0, 255).astype(np.uint8)
            
            return overlap_img
            
        except Exception as e:
            print(f"Create overlap visualization failed: {str(e)}")
            return None
    
if __name__ == "__main__":
               # 使用config.py中的默认值或用户提供的值
    model_path = r"devdeploy\weights\JinXiang_ZZ\fullmodel_best.onnx"
    input_folder = r"devdeploy\data\zz\ori_data"
    output_folder = r"devdeploy\data\zz\ori_data\test_Data"
    target_classes = [1, 2]
    device = "cuda"
    
    # 解析目标类别，如果没有提供则使用config.py中的默认值
    
    target_classes_list = [int(x) for x in target_classes]
    
    # 初始化分析服务
    service = SegmentationAnalyzerService(
        model_path=model_path,
        device=device
    )
    
    # 执行分析
    result = service.analyze_folder(
        input_folder=input_folder,
        output_folder=output_folder,
        target_classes=target_classes_list
    )
    from fastapi.responses import JSONResponse
    re=JSONResponse(content=result)
    print (re)